import numpy as np
import cv2
import onnxruntime as ort
import matplotlib.pyplot as plt

from cxf_utils import my_timer
import os
from ultralytics import YOLO
import torch
import time

if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    m = YOLO("wheel_seg_best.pt")
    m.to(device)

    t_s = time.time()
    img_source=plt.imread("./13.bmp")
    results = m(img_source)  # list of Results objects
    print(time.time() - t_s)
    masks = results[0].masks.data
    mask1 = np.uint8(np.array(masks)[0,:,:]*255)
    mask2 = np.uint8(np.array(masks)[1, :, :]*255)
    # kernel = np.ones((5, 5), dtype=np.int8)
    # mask1 = cv2.morphologyEx(mask1, cv2.MORPH_OPEN, kernel, iterations=2)  # 闭运算：先膨胀后腐蚀
    # mask2 = cv2.morphologyEx(mask2, cv2.MORPH_OPEN, kernel, iterations=2)  # 闭运算：先膨胀后腐蚀

    img_source = cv2.resize(img_source, (mask1.shape[1], mask1.shape[0]))


    contours, hierarchy = cv2.findContours(mask1,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
    contours_num_list=[]
    for i in range(len(contours)):
        contours_num_list.append(contours[i].shape[0])
    ellipse = cv2.fitEllipse(contours[np.argmax(np.array(contours_num_list))])
    #cv2.ellipse(img_source, ellipse, (0, 0, 255), 1, cv2.LINE_AA)
    x, y = ellipse[0]
    r=np.mean(ellipse[1])/2
    cv2.circle(img_source, (int(x), int(y)), int(r), (255, 0, 0), 1)

    contours, hierarchy = cv2.findContours(mask2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours_num_list = []
    for i in range(len(contours)):
        contours_num_list.append(contours[i].shape[0])
    ellipse = cv2.fitEllipse(contours[np.argmax(np.array(contours_num_list))])
    # cv2.ellipse(img_source, ellipse, (0, 0, 255), 1, cv2.LINE_AA)
    x, y = ellipse[0]
    r = np.mean(ellipse[1]) / 2
    cv2.circle(img_source, (int(x), int(y)), int(r), (0, 0, 255), 1)

    plt.imshow(img_source)
    plt.show()

    plt.imshow(mask1)
    plt.show()

    plt.imshow(mask2)
    plt.show()

    # mask1=np.zeros_like(mask1)
    # mask1_a = contours[0]
    # for i in range(mask1_a.shape[0]):
    #     mask1[mask1_a[i,0,1],mask1_a[i,0,0]]=255
    #
    # contours, hierarchy = cv2.findContours(mask2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    # mask2=np.zeros_like(mask2)
    # mask2_a = contours[0]
    # for i in range(mask2_a.shape[0]):
    #     mask2[mask2_a[i,0,1],mask2_a[i,0,0]]=255
    #
    # img_out = cv2.HoughCircles(mask1, cv2.HOUGH_GRADIENT, 1, 10, param1=200, param2=30, minRadius=0,
    #                            maxRadius=500)
    # img_source=cv2.resize(img_source,(mask1.shape[1],mask1.shape[0]))
    # for i in img_out[0, :]:
    #     cv2.circle(img_source, (int(i[0]), int(i[1])), int(i[2]), (255, 0, 0), 8)  # 绘制检测到的圆
    # plt.imshow(img_source, cmap="gray")
    # plt.show()
    #
    # plt.imshow(mask1)
    # plt.show()
    #
    # plt.imshow(mask2)
    # plt.show()

